package com.zh.stargraph2.core.service.impl;

import com.zh.stargraph2.core.service.OllamaService;
import com.zh.stargraph2.ollama.client.api.OllamaApi;
import com.zh.stargraph2.ollama.client.pojo.OllamaChatRequest;
import com.zh.stargraph2.ollama.client.pojo.OllamaChatRespone;
import com.zh.stargraph2.ollama.client.pojo.OllamaMessage;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import retrofit2.Call;
import retrofit2.Response;

import java.io.IOException;
import java.util.List;

@Service
public class OllamaServiceImpl implements OllamaService {

    @Autowired
    OllamaApi ollamaApi;


    @Override
    public String translate(String prompt) {
        OllamaMessage message = new OllamaMessage();
        message.setContent(String.format("请把下面文本翻译成英文：\n%s", prompt));
        message.setRole("user");
        OllamaChatRequest ollamaChatRequest = new OllamaChatRequest();
        ollamaChatRequest.setModel("qwen2:0.5b");
        ollamaChatRequest.setMessages(List.of(message));
        Call<OllamaChatRespone> chat = ollamaApi.chat(ollamaChatRequest);
        Response<OllamaChatRespone> execute = null;
        try {
            execute = chat.execute();
            if (execute.isSuccessful()) {
                OllamaChatRespone body = execute.body();
                System.out.println(body);
                return body.getMessage().getContent();
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
        return prompt;
    }
}
